import torch
import torch.nn as nn
from encoder import ResNet
from bridge_network import TF_model
from decoder_1212_revise import Decoder


class Mymodel(nn.Module):
    def __init__(self):
        super(Mymodel, self).__init__()
        self.encoder = ResNet()
        self.mctrans = TF_model()
        self.decoder = Decoder()

    def forward(self, input):
        feats = self.encoder(input)
        feats = self.mctrans(feats)
        output = self.decoder(feats)

        return output


if __name__ == '__main__':
    img = torch.Tensor(1, 1, 384, 384).cuda()
    net = Mymodel().cuda()
    out = net(img)
    print("out.shape",out.shape)
